memoryCache = {}

def cacheInMemory(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        global memoryCache
        key = func.__name__
        if key not in memoryCache:
            memoryCache[key] = func(*args, **kwargs)
        return memoryCache[key]
    return wrapper

class ImportDll:
    def __init__(self, name, nameFuncs=None, ordFuncs=None):
        self.name = name
        self.nameFuncs = nameFuncs or []
        self.ordFuncs = ordFuncs or []

    def addNameFunc(self, funcName):
        self.nameFuncs.append(funcName)

    def addOrdFunc(self, ordFunc):
        self.ordFuncs.append(ordFunc)

    def __str__(self):
        return f"""
            name: {self.name}
            nameFuncs: {self.nameFuncs}
            ordFuncs: {self.ordFuncs}
        """


class DllAnalyzer:

    @staticmethod
    @cacheInMemory
    def readDLLsImportTable(libPaths):
        # 结果
        dllMaps = {}
        # 可执行文件地址
        execFiles = []
        fileFilter = lambda name: name.lower().endswith(".exe") or name.lower().endswith(".dll")
        for libPath in libPaths:
            if os.path.isfile(libPath):
                if fileFilter(libPath):
                    execFiles.append(libPath)
            elif os.path.isdir(libPath):
                for root, _, files in os.walk(libPath):
                    for file in files:
                        fullPath = root + os.path.sep + file
                        if fileFilter(fullPath):
                            execFiles.append(fullPath)
            else:
                pass
        for file in execFiles:
            dllMaps[os.path.basename(file).lower()] = DllAnalyzer.readDLLImportTable(file)
        return dllMaps

    @staticmethod
    def readDLLImportTable(dllPath):
        importDlls = {}
        try:
            # 加载 DLL 文件
            pe = pefile.PE(dllPath)
            # 检查是否存在导入表
            if hasattr(pe, 'DIRECTORY_ENTRY_IMPORT'):
                for entry in pe.DIRECTORY_ENTRY_IMPORT:
                    dllName = entry.dll.decode("utf-8").lower()
                    importDlls[dllName] = ImportDll(dllName)
                    for imp in entry.imports:
                        if imp.name:
                            importDlls[dllName].addNameFunc(imp.name.decode('utf-8').lower())
                        else:
                            importDlls[dllName].addOrdFunc(imp.ordinal)
            else:
                print("此 DLL 文件没有导入表。")
        except pefile.PEFormatError:
            print("错误: 文件不是有效的 PE 格式。")
        except FileNotFoundError:
            print("错误: 文件未找到。")
        except Exception as e:
            print(f"发生未知错误: {e}")
        return importDlls


class DllAnalyzerVisitor:
    @staticmethod
    def printDllImports(libPaths):
        dllMaps = DllAnalyzer.readDLLsImportTable(libPaths)
        for dll, value in dllMaps.items():
            print(dll)
            for importedDll in value.keys():
                print("\t" + importedDll)

    @staticmethod
    def generateMermaid(entry, libPaths):
        dllMaps = DllAnalyzer.readDLLsImportTable(libPaths)
        if entry not in dllMaps:
            print("无法找到入口")
            return ""
        visited = set()
        mermaid = "```mermaid\ngraph TB\n"
        queue = Queue()
        queue.put(entry)
        # 初始化
        while not queue.empty():
            father = queue.get()
            visited.add(father)
            if father in dllMaps:
                for child in dllMaps.get(father).keys():
                    if child not in visited:
                        queue.put(child)
                    mermaid += f"""{father}-->{child}\n"""
        mermaid += "```\n"
        return mermaid

    @staticmethod
    def whoImportDLL(dll, libPaths):
        dllMaps = DllAnalyzer.readDLLsImportTable(libPaths)
        for parent, value in dllMaps.items():
            if dll in value:
                print(parent)

    @staticmethod
    def whatImport(dll, libPaths):
        dllMaps = DllAnalyzer.readDLLsImportTable(libPaths)
        if dll in dllMaps:
            for child in dllMaps.get(dll):
                print(child)

def Console(consolePaths):
    consolePaths = consolePaths
    cmds = {
        "whoImport": lambda item: DllAnalyzerVisitor.whoImportDLL(item, consolePaths),
        "refresh": lambda _: memoryCache.clear(),
        "printCache": lambda _: print(memoryCache),
        "generateMermaid": lambda item: DllAnalyzerVisitor.generateMermaid(item, consolePaths),
        "whatImport": lambda item: DllAnalyzerVisitor.whatImport(item, consolePaths),
        "help": lambda _: print(list(cmds.keys()))
    }
    content = input(">> ")
    while content not in ["quit", "exit"]:
        arr = content.split(" ")
        cmd = arr[0]
        if cmd in cmds:
            cmds[cmd](" ".join(arr[1:]))
        else:
            print("未知命令")
        content = input(">> ")